#include <iostream>
#include <string>
#include <chrono>
#include <iomanip>
#include <thread>
#include "video_processor.h"
#include "ball_tracker.h"
#include "utils.h"

void printUsage(const char* program_name) {
    std::cout << "Usage: " << program_name << " <engine_path> <input_video> <output_dir> [max_frames]" << std::endl;
    std::cout << "Arguments:" << std::endl;
    std::cout << "  engine_path   : Path to TensorRT engine file" << std::endl;
    std::cout << "  input_video   : Path to input video file" << std::endl;
    std::cout << "  output_dir    : Directory for output files" << std::endl;
    std::cout << "  max_frames    : Maximum number of frames to process (optional)" << std::endl;
    std::cout << std::endl;
    std::cout << "Example: " << program_name << " tracknetv3.engine input.mp4 ./output 1000" << std::endl;
}

int main(int argc, char* argv[]) {
    if (argc < 4) {
        printUsage(argv[0]);
        return -1;
    }
    
    std::string engine_path = argv[1];
    std::string input_video = argv[2];
    std::string output_dir = argv[3];
    int max_frames = (argc > 4) ? std::stoi(argv[4]) : 0;
    
    // 检查输入文件是否存在
    if (!fileExists(engine_path)) {
        std::cerr << "Error: Engine file not found: " << engine_path << std::endl;
        return -1;
    }
    
    if (!fileExists(input_video)) {
        std::cerr << "Error: Input video not found: " << input_video << std::endl;
        return -1;
    }
    
    // 创建输出目录
    if (!directoryExists(output_dir)) {
        std::cout << "Creating output directory: " << output_dir << std::endl;
        if (!createDirectory(output_dir)) {
            std::cerr << "Error: Failed to create output directory: " << output_dir << std::endl;
            return -1;
        }
    }
    
    // 生成输出文件路径
    std::string video_name = getFileNameWithoutExtension(input_video);
    std::string output_video = output_dir + "/" + video_name + "_tracked.mp4";
    std::string output_csv = output_dir + "/" + video_name + "_track.csv";
    
    std::cout << "=== TrackNetV3 Badminton Ball Tracking (GPU Accelerated) ===" << std::endl;
    std::cout << "Engine file: " << engine_path << std::endl;
    std::cout << "Input video: " << input_video << std::endl;
    std::cout << "Output video: " << output_video << std::endl;
    std::cout << "Output CSV: " << output_csv << std::endl;
    std::cout << "Max frames: " << (max_frames > 0 ? std::to_string(max_frames) : "All") << std::endl;
    std::cout << "===========================================================" << std::endl;
    
    try {
        // 初始化视频处理器
        VideoProcessor video_processor;
        std::cout << "Opening video file: " << input_video << std::endl;
        if (!video_processor.open(input_video)) {
            std::cerr << "Failed to open video: " << input_video << std::endl;
            return -1;
        }
        
        const VideoInfo& video_info = video_processor.getVideoInfo();
        
        // 尝试初始化硬件解码器
        std::cout << "Initializing hardware decoder..." << std::endl;
        bool hardware_decoder_available = video_processor.initHardwareDecoder();
        if (!hardware_decoder_available) {
            assert(false);
            std::cout << "Hardware decoder not available, using software decoder" << std::endl;
        } else {
            std::cout << "Hardware decoder initialized successfully" << std::endl;
        }
    
        
        // 初始化球跟踪器
        BallTracker ball_tracker;
        std::cout << "Initializing ball tracker with engine: " << engine_path << std::endl;
        if (!ball_tracker.init(engine_path)) {
            std::cerr << "Failed to initialize ball tracker" << std::endl;
            return -1;
        }
        std::cout << "Ball tracker initialized successfully" << std::endl;
        
        // 初始化GPU缓冲区
        std::cout << "Initializing GPU buffers for resolution: " 
                  << video_info.width << "x" << video_info.height << std::endl;
        
        if (!ball_tracker.initGPUBuffers(video_info.width, video_info.height)) {
            std::cerr << "Failed to initialize GPU buffers" << std::endl;
            return -1;
        }
        std::cout << "GPU buffers initialized successfully" << std::endl;
        
        ball_tracker.setOutputPaths(output_video, output_csv);
        
        std::cout << "Starting GPU-accelerated video processing..." << std::endl;
        auto total_start = std::chrono::high_resolution_clock::now();
        
        int processed_frames = 0;
        int detected_frames = 0;
        int consecutive_failures = 0;
        const int MAX_CONSECUTIVE_FAILURES = 10;
        
        // 使用GPU处理流程
        while (true) {
            if (max_frames > 0 && processed_frames >= max_frames) {
                std::cout << "Reached maximum frame limit: " << max_frames << std::endl;
                break;
            }
            
            try {
                bool success = false;
                if (hardware_decoder_available) {
                    success = ball_tracker.processFrameGPU(video_processor);
                } else {
                    assert(false);
                    success = ball_tracker.processFrameSoftware(video_processor); 
                }
                
                if (!success) {
                    assert(false);
                    consecutive_failures++;
                    std::cout << "Frame processing failed at frame " << processed_frames 
                              << " (consecutive failures: " << consecutive_failures << ")" << std::endl;
                    
                    if (consecutive_failures >= MAX_CONSECUTIVE_FAILURES) {
                        std::cout << "Too many consecutive failures, stopping processing" << std::endl;
                        break;
                    }
                    
                    // 短暂暂停后继续尝试
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                    continue;
                }
                
                // 重置连续失败计数
                consecutive_failures = 0;
                
                // 检查是否检测到球
                const BallPosition& ball_pos = ball_tracker.getCurrentPosition();
                if (ball_pos.confidence > 0.3f) {
                    detected_frames++;
                }
                
                processed_frames++;

                if (processed_frames % 50 == 0) {
                    std::cout << "=== Processing Status ===" << std::endl;
                    std::cout << "Processed frames: " << processed_frames << std::endl;
                    std::cout << "Current frame count: " << ball_tracker.getCurrentPosition().frame_id << std::endl;
                    
                    // 检查文件状态
                    std::ifstream csv_test(ball_tracker.getCSVPath());
                    if (csv_test.is_open()) {
                        std::cout << "CSV file is accessible" << std::endl;
                        csv_test.close();
                    } else {
                        std::cout << "CSV file may have issues" << std::endl;
                    }
                }
                
                // 显示进度
                if (processed_frames % 100 == 0) {
                    std::cout << "Processed " << processed_frames << " frames..." 
                              << " (detected: " << detected_frames << ")" << std::endl;
                    
                    // 显示处理速度
                    auto current_time = std::chrono::high_resolution_clock::now();
                    auto elapsed_seconds = std::chrono::duration_cast<std::chrono::seconds>(
                        current_time - total_start).count();
                    
                    if (elapsed_seconds > 0) {
                        double fps = static_cast<double>(processed_frames) / elapsed_seconds;
                        std::cout << "Current processing speed: " << std::fixed << std::setprecision(2) 
                                  << fps << " FPS" << std::endl;
                    }
                }
                
            } catch (const std::exception& e) {
                std::cerr << "Exception during frame processing at frame " << processed_frames 
                          << ": " << e.what() << std::endl;
                consecutive_failures++;
                
                if (consecutive_failures >= MAX_CONSECUTIVE_FAILURES) {
                    std::cout << "Too many consecutive exceptions, stopping processing" << std::endl;
                    break;
                }
            } catch (...) {
                std::cerr << "Unknown exception during frame processing at frame " << processed_frames << std::endl;
                consecutive_failures++;
                
                if (consecutive_failures >= MAX_CONSECUTIVE_FAILURES) {
                    std::cout << "Too many consecutive unknown exceptions, stopping processing" << std::endl;
                    break;
                }
            }
        }
        
        auto total_end = std::chrono::high_resolution_clock::now();
        auto total_duration = std::chrono::duration_cast<std::chrono::seconds>(total_end - total_start);
        
        // 释放资源
        std::cout << "Finalizing outputs..." << std::endl;
        ball_tracker.writeOutputs();
        video_processor.close();
        
        // 输出统计信息
        std::cout << "===========================================" << std::endl;
        std::cout << "GPU-accelerated processing completed!" << std::endl;
        std::cout << "Total frames processed: " << processed_frames << std::endl;
        std::cout << "Frames with ball detected: " << detected_frames << std::endl;
        
        if (processed_frames > 0) {
            double detection_rate = (static_cast<double>(detected_frames) / processed_frames) * 100.0;
            std::cout << "Detection rate: " << std::fixed << std::setprecision(2) 
                      << detection_rate << "%" << std::endl;
            
            std::cout << "Total processing time: " << total_duration.count() << " seconds" << std::endl;
            
            if (total_duration.count() > 0) {
                double avg_fps = static_cast<double>(processed_frames) / total_duration.count();
                std::cout << "Average FPS: " << std::fixed << std::setprecision(2) << avg_fps << std::endl;
            }
        } else {
            std::cout << "No frames were successfully processed" << std::endl;
        }
        
        std::cout << "Output files:" << std::endl;
        std::cout << "  - Video: " << output_video << std::endl;
        std::cout << "  - CSV: " << output_csv << std::endl;
        std::cout << "Output directory: " << output_dir << std::endl;
        std::cout << "===========================================" << std::endl;
        
    } catch (const std::exception& e) {
        std::cerr << "Exception occurred in main: " << e.what() << std::endl;
        return -1;
    } catch (...) {
        std::cerr << "Unknown exception occurred in main" << std::endl;
        return -1;
    }
    
    std::cout << "Program completed successfully" << std::endl;
    return 0;
}
